#!/bin/bash
# Copyright (c) 2018-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

############################################################################
## This script generates the model repository needed by some of the
## tritonserver CI tests. Generating these models requires using
## the TensorFlow and PyTorch containers.
##
## 1. Update TENSORRT_IMAGE, PYTORCH_IMAGE and TENSORFLOW_IMAGE to
## match what is being used by the tritonserver release being
## tested.
##
## 2. Set CUDA_DEVICE to the ID of the CUDA device present on the
## local system that you want to target for the generated models.
##
## 3. Run this script to create /tmp/qa_model_repository,
## /tmp/qa_variable_model_repository, /tmp/qa_shapetensor_model_repository
## /tmp/qa_identity_model_repository, /tmp/qa_identity_big_model_repository
## /tmp/qa_reshape_model_repository, /tmp/qa_noshape_model_repository,
## /tmp/qa_sequence_model_repository, /tmp/qa_ensemble_model_repository
## /tmp/qa_dyna_sequence_model_repository, and
## /tmp/qa_variable_sequence_model_repository directories containing
## all the models needed for CI testing.
##
############################################################################

TRITON_VERSION=20.11

# ONNX. Use ONNX_OPSET 0 to use the default for ONNX version
ONNX_VERSION=1.6.0
ONNX_OPSET=0

ONNX_IMAGE=${ONNX_IMAGE:=ubuntu:18.04}
PYTORCH_IMAGE=${PYTORCH_IMAGE:=nvcr.io/nvidia/pytorch:$TRITON_VERSION-py3}
TENSORFLOW_IMAGE=${TENSORFLOW_IMAGE:=nvcr.io/nvidia/tensorflow:$TRITON_VERSION-tf1-py3}
TENSORRT_IMAGE=${TENSORRT_IMAGE:=nvcr.io/nvidia/tensorrt:$TRITON_VERSION-py3}
CUDA_DEVICE=0

###
HOST_SRCDIR=/tmp/gen_srcdir
HOST_DESTDIR=/tmp/$TRITON_VERSION/qa_model_repository
HOST_VARDESTDIR=/tmp/$TRITON_VERSION/qa_variable_model_repository
HOST_IDENTITYDESTDIR=/tmp/$TRITON_VERSION/qa_identity_model_repository
HOST_IDENTITYBIGDESTDIR=/tmp/$TRITON_VERSION/qa_identity_big_model_repository
HOST_SHAPEDESTDIR=/tmp/$TRITON_VERSION/qa_shapetensor_model_repository
HOST_RESHAPEDESTDIR=/tmp/$TRITON_VERSION/qa_reshape_model_repository
HOST_SEQDESTDIR=/tmp/$TRITON_VERSION/qa_sequence_model_repository
HOST_DYNASEQDESTDIR=/tmp/$TRITON_VERSION/qa_dyna_sequence_model_repository
HOST_VARSEQDESTDIR=/tmp/$TRITON_VERSION/qa_variable_sequence_model_repository
HOST_ENSEMBLEDESTDIR=/tmp/$TRITON_VERSION/qa_ensemble_model_repository
HOST_NOSHAPEDESTDIR=/tmp/$TRITON_VERSION/qa_noshape_model_repository
HOST_PLGDESTDIR=/tmp/$TRITON_VERSION/qa_trt_plugin_model_repository
HOST_RAGGEDDESTDIR=/tmp/$TRITON_VERSION/qa_ragged_model_repository
HOST_FORMATDESTDIR=/tmp/$TRITON_VERSION/qa_trt_format_model_repository

rm -fr $HOST_SRCDIR $HOST_DESTDIR $HOST_VARDESTDIR
rm -fr $HOST_IDENTITYDESTDIR $HOST_IDENTITYBIGDESTDIR $HOST_SHAPEDESTDIR
rm -fr $HOST_SEQDESTDIR $HOST_DYNASEQDESTDIR $HOST_VARSEQDESTDIR
rm -fr $HOST_ENSEMBLEDESTDIR $HOST_NOSHAPEDESTDIR $HOST_RESHAPEDESTDIR
rm -fr $HOST_PLGDESTDIR $HOST_RAGGEDDESTDIR $HOST_FORMATDESTDIR
mkdir -p $HOST_SRCDIR
mkdir -p $HOST_DESTDIR
mkdir -p $HOST_VARDESTDIR
mkdir -p $HOST_IDENTITYDESTDIR
mkdir -p $HOST_IDENTITYBIGDESTDIR
mkdir -p $HOST_SHAPEDESTDIR
mkdir -p $HOST_RESHAPEDESTDIR
mkdir -p $HOST_SEQDESTDIR
mkdir -p $HOST_DYNASEQDESTDIR
mkdir -p $HOST_VARSEQDESTDIR
mkdir -p $HOST_ENSEMBLEDESTDIR
mkdir -p $HOST_NOSHAPEDESTDIR
mkdir -p $HOST_PLGDESTDIR
mkdir -p $HOST_RAGGEDDESTDIR
mkdir -p $HOST_FORMATDESTDIR

# Since the models required by ensemble models may not be available
# at this point, storing ensemble models separately so that other qa directories
# are proper model repositories
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_variable_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_identity_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_reshape_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_sequence_model_repository
mkdir -p $HOST_ENSEMBLEDESTDIR/qa_variable_sequence_model_repository

rm -fr $HOST_SRCDIR
mkdir -p $HOST_SRCDIR

cp ./gen_qa_models.py $HOST_SRCDIR/.
cp ./gen_qa_identity_models.py $HOST_SRCDIR/.
cp ./gen_qa_reshape_models.py $HOST_SRCDIR/.
cp ./gen_qa_noshape_models.py $HOST_SRCDIR/.
cp ./gen_qa_sequence_models.py $HOST_SRCDIR/.
cp ./gen_qa_dyna_sequence_models.py $HOST_SRCDIR/.
cp ./gen_ensemble_model_utils.py $HOST_SRCDIR/.
cp ./gen_qa_trt_plugin_models.py $HOST_SRCDIR/.
cp ./gen_qa_trt_format_models.py $HOST_SRCDIR/.
cp ./gen_qa_ragged_models.py $HOST_SRCDIR/.
cp ./test_util.py $HOST_SRCDIR/.

ONNXSCRIPT=onnx_gen.cmds
TORCHSCRIPT=torch_gen.cmds
TFSCRIPT=tf_gen.cmds
TRTSCRIPT=trt_gen.cmds

SRCDIR=/tmp/src
DESTDIR=/tmp/models
VARDESTDIR=/tmp/varmodels
IDENTITYDESTDIR=/tmp/zeromodels
IDENTITYBIGDESTDIR=/tmp/zerobigmodels
SHAPEDESTDIR=/tmp/shapetensormodels
RESHAPEDESTDIR=/tmp/reshapemodels
SEQDESTDIR=/tmp/seqmodels
DYNASEQDESTDIR=/tmp/dynaseqmodels
VARSEQDESTDIR=/tmp/varseqmodels
ENSEMBLEDESTDIR=/tmp/ensemblemodels
NOSHAPEDESTDIR=/tmp/noshapemodels
PLGDESTDIR=/tmp/pluginmodels
RAGGEDDESTDIR=/tmp/raggedmodels
FORMATDESTDIR=/tmp/formatmodels

# ONNX
cat >$HOST_SRCDIR/$ONNXSCRIPT <<EOF
#!/bin/bash
set -e
apt-get update && apt-get install -y --no-install-recommends python3 python3-pip
pip3 install --upgrade onnx==${ONNX_VERSION}

python3 $SRCDIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$ONNXSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $ONNX_IMAGE
docker run --gpus device=$CUDA_DEVICE --rm --entrypoint $SRCDIR/$ONNXSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       $ONNX_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# PyTorch
cat >$HOST_SRCDIR/$TORCHSCRIPT <<EOF
#!/bin/bash
set -e
python3 $SRCDIR/gen_qa_models.py --libtorch --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --libtorch --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --libtorch --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --libtorch --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --libtorch --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --libtorch --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --libtorch --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$TORCHSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $PYTORCH_IMAGE
docker run --gpus device=$CUDA_DEVICE --rm --entrypoint $SRCDIR/$TORCHSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       $PYTORCH_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# Tensorflow
cat >$HOST_SRCDIR/$TFSCRIPT <<EOF
#!/bin/bash
set -e
python3 $SRCDIR/gen_qa_models.py --graphdef --savedmodel --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --graphdef --savedmodel --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --graphdef --savedmodel --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --graphdef --savedmodel --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --graphdef --savedmodel --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --graphdef --savedmodel --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --graphdef --savedmodel --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_noshape_models.py --savedmodel --models_dir=$NOSHAPEDESTDIR
chmod -R 777 $NOSHAPEDESTDIR
python3 $SRCDIR/gen_qa_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_model_repository
python3 $SRCDIR/gen_qa_models.py --ensemble --variable --models_dir=$ENSEMBLEDESTDIR/qa_variable_model_repository
python3 $SRCDIR/gen_qa_reshape_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_reshape_model_repository
python3 $SRCDIR/gen_qa_identity_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_identity_model_repository
python3 $SRCDIR/gen_qa_sequence_models.py --ensemble --models_dir=$ENSEMBLEDESTDIR/qa_sequence_model_repository
python3 $SRCDIR/gen_qa_sequence_models.py --ensemble --variable --models_dir=$ENSEMBLEDESTDIR/qa_variable_sequence_model_repository
chmod -R 777 $ENSEMBLEDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$TFSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $TENSORFLOW_IMAGE
docker run --gpus device=$CUDA_DEVICE --rm --entrypoint $SRCDIR/$TFSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       --mount type=bind,source=$HOST_NOSHAPEDESTDIR,target=$NOSHAPEDESTDIR \
       --mount type=bind,source=$HOST_ENSEMBLEDESTDIR,target=$ENSEMBLEDESTDIR \
       $TENSORFLOW_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# TensorRT
cat >$HOST_SRCDIR/$TRTSCRIPT <<EOF
#!/bin/bash
set -e
export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
# Models using shape tensor i/o
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
chmod -R 777 $SHAPEDESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-big --models_dir=$IDENTITYBIGDESTDIR
chmod -R 777 $IDENTITYBIGDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --tensorrt --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_ragged_models.py --tensorrt --models_dir=$RAGGEDDESTDIR
chmod -R 777 $RAGGEDDESTDIR
python3 $SRCDIR/gen_qa_trt_format_models.py --models_dir=$FORMATDESTDIR
chmod -R 777 $FORMATDESTDIR
# make shared library for custom clip plugin
(cd /workspace/tensorrt/samples/python/uff_custom_plugin && cmake . && make && \
cp libclipplugin.so $PLGDESTDIR/.)
LD_PRELOAD=$PLGDESTDIR/libclipplugin.so python3 $SRCDIR/gen_qa_trt_plugin_models.py --models_dir=$PLGDESTDIR
chmod -R 777 $PLGDESTDIR
EOF

chmod a+x $HOST_SRCDIR/$TRTSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker pull $TENSORRT_IMAGE
docker run --gpus device=$CUDA_DEVICE --rm --entrypoint $SRCDIR/$TRTSCRIPT \
       --mount type=bind,source=$HOST_SRCDIR,target=$SRCDIR \
       --mount type=bind,source=$HOST_DESTDIR,target=$DESTDIR \
       --mount type=bind,source=$HOST_VARDESTDIR,target=$VARDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYDESTDIR,target=$IDENTITYDESTDIR \
       --mount type=bind,source=$HOST_IDENTITYBIGDESTDIR,target=$IDENTITYBIGDESTDIR \
       --mount type=bind,source=$HOST_SHAPEDESTDIR,target=$SHAPEDESTDIR \
       --mount type=bind,source=$HOST_RESHAPEDESTDIR,target=$RESHAPEDESTDIR \
       --mount type=bind,source=$HOST_SEQDESTDIR,target=$SEQDESTDIR \
       --mount type=bind,source=$HOST_DYNASEQDESTDIR,target=$DYNASEQDESTDIR \
       --mount type=bind,source=$HOST_VARSEQDESTDIR,target=$VARSEQDESTDIR \
       --mount type=bind,source=$HOST_PLGDESTDIR,target=$PLGDESTDIR \
       --mount type=bind,source=$HOST_RAGGEDDESTDIR,target=$RAGGEDDESTDIR \
       --mount type=bind,source=$HOST_FORMATDESTDIR,target=$FORMATDESTDIR \
       $TENSORRT_IMAGE
if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi
